import numpy as np


def omega_s(h, omega_0, alpha_c, v_mc, phi_s, energy, e):
    """
    计算无束流负载时的同步振荡角频率
    :param h: 谐波数
    :param omega_0: 回旋角频率
    :param alpha_c: 动量紧缩因子
    :param v_mc: 主腔腔压
    :param phi_s: 主腔相位
    :param energy: 粒子能量，单位和电荷单位匹配就行
    :param e: 电荷绝对值
    """
    return omega_0 * np.sqrt(h * alpha_c * v_mc * np.abs(np.cos(phi_s)) / (2 * np.pi * energy / e))


def omega_positive(omega_0, omega_complex, p: int, m: int, mu: int):
    """
    计算
    :param omega_0: 回旋角频率
    :param omega_complex: 复振荡角频率，实部代表相干角频率，虚部代表不稳定性增长率
    :param p: 整数
    :param m: 整数
    :param mu: 整数，取值[0,m-1]
    """
    return p * m * omega_0 + mu * omega_0 + omega_complex


def omega_negative(omega_0, omega_complex, p: int, m: int, mu: int):
    """
    计算
    :param omega_0: 回旋角频率
    :param omega_complex: 复振荡角频率，实部代表相干角频率，虚部代表不稳定性增长率
    :param p: 整数
    :param m: 整数
    :param mu: 整数，取值[0,m-1]
    """
    return p * m * omega_0 - mu * omega_0 - omega_complex


def z_op(omega: np, r_l, q_l, omega_r) -> np:
    """计算开环阻抗"""
    return r_l / (1 + 1j * q_l * (omega / omega_r - omega_r / omega))


def h_rf(omega: np, tau, k_p, phi, r_l) -> np:
    """rf反馈的传递函数"""
    return np.exp(-1j * tau * omega + 1j * phi) * k_p / r_l


def h_otfb(omega: np, g, k, t_g, t_0) -> np:
    """OTFB的传递函数"""
    return g / (1 - k * np.exp(-1j * omega * t_0)) * np.exp(-1j * omega * (t_0 - t_g))


# def z_otfb(omega: np, r_l, q_l, omega_r, tau, k_p, phi, g, k, t_g, t_0) -> np:
#     """
#     反馈后的总阻抗
#     :param omega: 角频率
#     :param r_l: 分路阻抗
#     :param q_l: 品质因子
#     :param omega_r: 谐振频率
#     :param tau: 直接反馈延迟
#     :param k_p: 直接反馈增益
#     :param phi: 直接反馈相移
#     :param g: OTFB怎用
#     :param k: OTFB带宽相关参数
#     :param t_g: OTFB延迟
#     :param t_0: 回旋时间
#     """
#     z_ol = z_op(omega, r_l, q_l, omega_r)
#     rf_fb = h_rf(omega, tau, k_p, phi, r_l)
#     otfb = h_otfb(omega, g, k, t_g, t_0)
#     return z_ol / (1 + z_ol * rf_fb * (1 + otfb))


class MyObj:
    def __init__(self, omega_0, omega_complex, r_l, q_l, omega_r, tau, k_p, phi, g, k, t_g, t_0):
        """
        :param omega_0: 回旋角频率
        :param omega_complex: 复振荡角频率，实部代表相干角频率，虚部代表不稳定性增长率，计算的时候取无束流负载时的同步振荡角频率
        :param r_l: 分路阻抗
        :param q_l: 品质因子
        :param omega_r: 谐振频率
        :param tau: 直接反馈延迟
        :param k_p: 直接反馈增益
        :param phi: 直接反馈相移 / 弧度
        :param g: OTFB怎用
        :param k: OTFB带宽相关参数
        :param t_g: OTFB延迟/回旋时间
        :param t_0: 回旋时间
        """
        self.omega_0 = omega_0
        self.omega_complex = omega_complex
        self.r_l = r_l
        self.q_l = q_l
        self.omega_r = omega_r
        self.tau = tau
        self.k_p = k_p
        self.phi = phi
        self.g = g
        self.k = k
        self.t_g = t_g * t_0
        self.t_0 = t_0

    def z_otfb(self, omega):
        """
        :param omega: 计算的角频率点
        :return:
        """
        z_ol = z_op(omega=omega, r_l=self.r_l, q_l=self.q_l, omega_r=self.omega_r)
        rf_fb = h_rf(omega=omega, tau=self.tau, k_p=self.k_p, phi=self.phi, r_l=self.r_l)
        otfb = h_otfb(omega=omega, g=self.g, k=self.k, t_g=self.t_g, t_0=self.t_0)
        return z_ol / (1 + z_ol * rf_fb * (1 + otfb))

    def p_re(self, p, m, mu):
        omega_p = omega_positive(omega_0=self.omega_0, omega_complex=self.omega_complex, p=p, m=m, mu=mu)
        omega_n = omega_negative(omega_0=self.omega_0, omega_complex=self.omega_complex, p=p, m=m, mu=mu)
        if isinstance(p, int):
            return np.abs(np.real(omega_p * self.z_otfb(omega_p) - omega_n * self.z_otfb(omega_n)))
        else:
            return sum(np.abs(
                [np.real(omega_p[i] * self.z_otfb(omega_p[i]) - omega_n[i] * self.z_otfb(omega_n[i])) for i in
                 range(len(p))]))

    def p_im(self, p, m, mu):
        omega_p = omega_positive(omega_0=self.omega_0, omega_complex=self.omega_complex, p=p, m=m, mu=mu)
        omega_n = omega_negative(omega_0=self.omega_0, omega_complex=self.omega_complex, p=p, m=m, mu=mu)
        if isinstance(p, int):
            return np.abs(np.imag(self.z_otfb(omega_p) + self.z_otfb(omega_n)))
        else:
            return sum(np.abs([np.imag(self.z_otfb(omega_p[i]) + self.z_otfb(omega_n[i])) for i in range(len(p))]))

    def phase_margin(self, omega, minangle, maxangle):
        """相位裕度约束条件"""
        try:
            pop = len(self.k)
        except TypeError:
            pop = 1

        def calculate(omega0, r_l=self.r_l, q_l=self.q_l, omega_r=self.omega_r, tau=self.tau, k_p=self.k_p,
                      phi=self.phi, g=self.g, k=self.k, t_g=self.t_g, t_0=self.t_0):
            z_ol = z_op(omega=omega0, r_l=r_l, q_l=q_l, omega_r=omega_r)
            rf_fb = h_rf(omega=omega0, tau=tau, k_p=k_p, phi=phi, r_l=r_l)
            otfb = h_otfb(omega=omega0, g=g, k=k, t_g=t_g, t_0=t_0)
            return z_ol * rf_fb * (1 + otfb)

        def zero(interval, tolerance=1e-4, r_l=self.r_l, q_l=self.q_l, omega_r=self.omega_r, tau=self.tau, k_p=self.k_p,
                 phi=self.phi, g=self.g, k=self.k, t_g=self.t_g, t_0=self.t_0):
            liter = int(np.log2(abs(interval[1] - interval[0]) / tolerance)) + 1
            if (
                    calculate(
                        interval[0], r_l=r_l, q_l=q_l, omega_r=omega_r, tau=tau, k_p=k_p, phi=phi, g=g, k=k, t_g=t_g,
                        t_0=t_0) >
                    calculate(
                        interval[1], r_l=r_l, q_l=q_l, omega_r=omega_r, tau=tau, k_p=k_p, phi=phi, g=g, k=k, t_g=t_g,
                        t_0=t_0)
            ):
                interval[0], interval[1] = interval[1], interval[0]
            for i in range(liter):
                res = (interval[0] + interval[1]) / 2
                if calculate(
                        res, r_l=r_l, q_l=q_l, omega_r=omega_r, tau=tau, k_p=k_p, phi=phi, g=g, k=k, t_g=t_g, t_0=t_0
                ) > 0:
                    interval[1] = res
                if calculate(
                        res, r_l=r_l, q_l=q_l, omega_r=omega_r, tau=tau, k_p=k_p, phi=phi, g=g, k=k, t_g=t_g, t_0=t_0
                ) < 0:
                    interval[0] = res
                else:
                    break
            return (interval[0] + interval[1]) / 2

        temp0 = np.array([calculate(i) for i in omega])  # 计算所有值
        temp0 = (np.abs(temp0[:-1]) - 1) * (np.abs(temp0[1:]) - 1)  # 为找异号点做准备
        result1 = []
        result2 = []
        if pop > 1:
            for j in range(pop):
                r_l, q_l, omega_r, tau, k_p, phi, g, k, t_g, t_0 = self.r_l, self.q_l, self.omega_r, self.tau, self.k_p[
                    j], self.phi[j], self.g[j], self.k[j], self.t_g[j], self.t_0
# TODO:
                temp1 = np.array([i for i in range(len(temp0)) if temp0[i, j] < 0])
                if len(temp1) > 0:
                    result1.append(- min([np.angle(
                        calculate(temp1[i], r_l=r_l, q_l=q_l, omega_r=omega_r, tau=tau, k_p=k_p, phi=phi, g=g, k=k,
                                  t_g=t_g, t_0=t_0)) / np.pi * 180 - minangle for i in range(len(temp1))])
                                   )
                    result2.append(max([np.angle(
                        calculate(temp1[i], r_l=r_l, q_l=q_l, omega_r=omega_r, tau=tau, k_p=k_p, phi=phi, g=g, k=k,
                                  t_g=t_g, t_0=t_0)) / np.pi * 180 - maxangle for i in range(len(temp1))]))
                else:
                    result1.append(100)
                    result2.append(100)
            return np.array(result1), np.array(result2)

        else:
            temp1 = np.array([i for i in range(len(temp0)) if temp0[i] < 0])  # 找到异号点
            return min(np.angle([calculate(temp1[i]) for i in range(len(temp1))])) / np.pi * 180, max(
                np.angle([calculate(temp1[i]) for i in range(len(temp1))]) / np.pi * 180)
# TODO:
        #     temp1 = np.array([i for i in range(len(temp0)) if temp0[i, j] < 0])
        #     zeros = [zero([temp0[temp1[i]], temp0[temp1[i]] + 1]) for i in temp1]
        #     if len(temp1) > 0:
        #         result1.append(- min([np.angle(
        #             calculate(zeros[i], r_l=r_l, q_l=q_l, omega_r=omega_r, tau=tau, k_p=k_p, phi=phi, g=g, k=k,
        #                       t_g=t_g, t_0=t_0)) / np.pi * 180 - minangle for i in range(len(zeros))])
        #                        )
        #         result2.append(max([np.angle(
        #             calculate(zeros[i], r_l=r_l, q_l=q_l, omega_r=omega_r, tau=tau, k_p=k_p, phi=phi, g=g, k=k,
        #                       t_g=t_g, t_0=t_0)) / np.pi * 180 - maxangle for i in range(len(zeros))]))
        #     else:
        #         result1.append(100)
        #         result2.append(100)
        # return np.array(result1), np.array(result2)
        #
        # else:
        # temp1 = np.array([i for i in range(len(temp0)) if temp0[i] < 0])  # 找到异号点
        # zeros = [zero([temp0[temp1[i]], temp0[temp1[i]] + 1]) for i in temp1]
        # return min(np.angle([calculate(zeros[i]) for i in range(len(zeros))])) / np.pi * 180, max(
        #     np.angle([calculate(zeros[i]) for i in range(len(zeros))]) / np.pi * 180)

    def calculate(self, omega0):
        z_ol = z_op(omega=omega0, r_l=self.r_l, q_l=self.q_l, omega_r=self.omega_r)
        rf_fb = h_rf(omega=omega0, tau=self.tau, k_p=self.k_p, phi=self.phi, r_l=self.r_l)
        otfb = h_otfb(omega=omega0, g=self.g, k=self.k, t_g=self.t_g, t_0=self.t_0)
        return z_ol * rf_fb * (1 + otfb)
